/****************************************************************************
 * Copyright (C) 2009-2010 SciTouch LLC
 * 
 * This file is part of Indigo toolkit.
 * 
 * This file may be distributed and/or modified under the terms of the
 * GNU General Public License version 3 as published by the Free Software
 * Foundation and appearing in the file LICENSE.GPL included in the
 * packaging of this file.
 * 
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 ***************************************************************************/

#include <ctype.h>

#include "molecule/smiles_loader.h"

#include "base_cpp/array.h"
#include "base_cpp/tlscont.h"
#include "base_cpp/scanner.h"

#include "molecule/molecule.h"
#include "molecule/molecule_stereocenters.h"
#include "graph/graph_highlighting.h"

SmilesLoader::SmilesLoader (Scanner &scanner) : _scanner(scanner),
TL_CP_GET(_atoms),
TL_CP_GET(_bonds),
TL_CP_GET(_closures)
{
   reaction_atom_mapping = 0;
   highlighting = 0;
   inside_rsmiles = false;
}

void SmilesLoader::loadMolecule (Molecule &mol, bool query)
{
   mol.clear();

   if (query)
   {
      mol.createQueryAtoms();
      mol.createQueryBonds();
   }

   QS_DEF(Array<_CycleDesc>, cycles);
   QS_DEF(Array<int>, atom_stack);

   _atoms.clear();
   _bonds.clear();
   _closures.clear();
   cycles.clear();
   atom_stack.clear();
   if (highlighting != 0)
      highlighting->clear();

   bool first_atom = true;

   while (!_scanner.isEOF())
   {
      char c = _scanner.readChar();

      if (isspace(c))
         break;

      _BondDesc *bond = 0;

      if (!first_atom)
      {
         bool added_bond = false;

         while (isdigit(c) || c == '%')
         {
            int number;

            if (c == '%')
               number = _scanner.readIntFix(2);
            else
               number = c - '0';

            while (cycles.size() <= number)
               cycles.push().clear();

            if (cycles[number].beg >= 0)
            {
               Edge &closure = _closures.push();

               closure.beg = cycles[number].beg;
               closure.end = atom_stack.top();

               bond = &_bonds.push();
               bond->dir = 0;
               bond->query.clear();
               bond->beg = atom_stack.top();
               bond->end = cycles[number].beg;
               bond->type = -1; // will later become single or aromatic bond
               cycles[number].clear();
               added_bond = true;
               break;
            }
            else if (cycles[number].pending_bond >= 0)
            {
               bond = &_bonds[cycles[number].pending_bond];
               bond->end = atom_stack.top();
               cycles[number].clear();
               added_bond = true;
               break;
            }
            else
            {
               cycles[number].beg = atom_stack.top();
               cycles[number].pending_bond = -1;
            }
            c = _scanner.readChar();
         }

         if (added_bond)
            continue;
      }

      if (c == '.')
      {
         atom_stack.pop();
         first_atom = true;
         continue;
      }

      if (c == '(')
      {
         atom_stack.push(atom_stack.top());
         continue;
      }

      if (c == ')')
      {
         atom_stack.pop();
         continue;
      }

      if (!first_atom)
      {
         bond = &_bonds.push();
         bond->beg = atom_stack.top();
         bond->end = -1;
         bond->type = -1;
         bond->dir = 0;
         bond->query.clear();
      }

      if (bond != 0)
      {
         bool bond_symbol = true;

         if (c == '-')
            bond->type = BOND_SINGLE;
         else if (c == '/')
         {
             bond->type = BOND_SINGLE;
             bond->dir = 1;
         }
         else if (c == '\\')
         {
             bond->type = BOND_SINGLE;
             bond->dir = 2;
         }
         else if (c == '=')
            bond->type = BOND_DOUBLE;
         else if (c == '#')
            bond->type = BOND_TRIPLE;
         else if (c == ':')
            bond->type = BOND_AROMATIC;
         else if (c == '~')
         {
            if (!query)
               throw Error("'~' any bond is allowed only for queries");
            bond->type = 0;
            bond->query.type = QUERY_BOND_ANY;
         }
         else
            bond_symbol = false;

         if (bond_symbol)
         {
            c = _scanner.readChar();

            if (isdigit(c) || c == '%')
            {
               int number;

               if (c == '%')
                  number = _scanner.readIntFix(2);
               else
                  number = c - '0';

               if (number >= 0 && number < cycles.size() && cycles[number].beg >= 0)
               {
                  Edge &closure = _closures.push();

                  closure.beg = cycles[number].beg;
                  closure.end = bond->beg;

                  bond->end = cycles[number].beg;
                  cycles[number].clear();
                  continue;
               }
               else if (number >= 0 && number < cycles.size() && cycles[number].pending_bond >= 0)
               {
                  _BondDesc &pending_bond = _bonds[cycles[number].pending_bond];

                  if (bond->type != pending_bond.type)
                     throw Error("cycle %d: closing bond (%d) does not match pending bond (%d)",
                        number, bond->type, pending_bond.type);

                  // transfer data from closing bond to pending bond
                  if (bond->dir > 0)
                  {
                     if (bond->dir == pending_bond.dir)
                        throw Error("cycle %d: closing bond direction does not match pending bond direction",
                           number);

                     pending_bond.dir = 3 - bond->dir;
                  }

                  pending_bond.end = bond->beg;

                  // forget the closing bond
                  _bonds.pop();
                  continue;
               }
               else
               {
                  while (cycles.size() <= number)
                     cycles.push().clear();
                  cycles[number].pending_bond = _bonds.size() - 1;
                  cycles[number].beg = -1; // have it already in the bond
                  continue;
               }
            }
         }

         if (c == '@')
         {
            if (!query)
               throw Error("'@' ring bond is allowed only for queries");
            bond->query.topology = QUERY_BOND_TOPOLOGY_RING;
            c = _scanner.readChar();
         }
         else if (c == '!')
         {
            c = _scanner.readChar();
            if (c != '@')
               throw Error("expecting '@' after '!', got %c", c);

            if (!query)
               throw Error("'!@' chain bond is allowed only for queries");

            bond->query.topology = QUERY_BOND_TOPOLOGY_CHAIN;
            c = _scanner.readChar();
         }

      }

      _AtomDesc &atom = _atoms.push();

      bool in_brackets = false;

      if (!first_atom)
         atom_stack.pop();
      atom_stack.push(_atoms.size() - 1);
      first_atom = false;

      if (bond != 0)
         bond->end = _atoms.size() - 1;

      memset(&atom, 0, sizeof(_AtomDesc));

      atom.hydrogens = -1;

      if (c == '[')
      {
         in_brackets = true;
         c = _scanner.readChar();
      }

      if (in_brackets)
      {
         while (isdigit(c))
         {
            atom.isotope = atom.isotope * 10 + c - '0';
            c = _scanner.readChar();
         }
      }

      if (c == '*')
      {
         // will decide later whether it is a query atom or pseudo-atom
         atom.query.type = QUERY_ATOM_A; 
      }
      /*else if (c == 'a')
      {
         if (!query)
            throw Error("'a' aromatic atoms allowed only for queries");
         atom.aromatic = true;
         atom.query.type = QUERY_ATOM_A;
         atom.query.aromaticity = QUERY_ATOM_AROMATIC;
      }
      else if (c == 'A')
      {
         if (!query)
            throw Error("'A' aliphatic atoms allowed only for queries");
         atom.query.type = QUERY_ATOM_A;
         atom.query.aromaticity = QUERY_ATOM_ALIPHATIC;
      }*/
      else
      {
         if (!isalpha((unsigned)c))
            throw Error("expected alphabet character, got %c", c);

         if (islower(c))
         {
            c = toupper(c);
            atom.aromatic = 1;
         }

         if (c == 'B' && _scanner.lookNext() == 'r')
         {
            _scanner.skip(1);
            atom.label = ELEM_Br;
         }
         else if (c == 'C' && _scanner.lookNext() == 'l')
         {
            _scanner.skip(1);
            atom.label = ELEM_Cl;
         }
         else if (in_brackets && isalpha(_scanner.lookNext()) && islower(_scanner.lookNext()))
            atom.label = Element::fromTwoChars(c, _scanner.readChar());
         else
            atom.label = Element::fromChar(c);
      }
      
      if (in_brackets)
      {
         atom.hydrogens = 0;
         c = _scanner.readChar();
         if (c == '@')
         {
            atom.chirality = 1;
            c = _scanner.readChar();
            if (c == '@')
            {
               c = _scanner.readChar();
               atom.chirality = 2;
            }
         }
         if (c == 'H')
         {
            c = _scanner.readChar();
            atom.hydrogens = 1;
            if (isdigit(c))
            {
               atom.hydrogens = 0;
               while (isdigit(c))
               {
                  atom.hydrogens = atom.hydrogens * 10 + c - '0';
                  c = _scanner.readChar();
               }
            }
         }
         if (c == '+' || c == '-')
         {
            char cstart = c;
            bool negative = (c == '-');

            c = _scanner.readChar();
            atom.charge = 1;
            if (isdigit(c))
            {
               atom.charge = 0;
               while (isdigit(c))
               {
                  atom.charge = atom.charge * 10 + c - '0';
                  c = _scanner.readChar();
               }
            }
            else while (c == cstart)
            {
               atom.charge++;
               c = _scanner.readChar();
            }
            if (negative)
               atom.charge = -atom.charge;
         }
         if (c == ':')
         {
            atom.aam = _scanner.readUnsigned();
            c = _scanner.readChar();
         }
         if (c != ']')
            throw Error("expected ], got %c", c);
      }
   }

   int i;

   for (i = 0; i < _atoms.size(); i++)
   {
      mol.addAtom(_atoms[i].label);

      Atom &atom = mol.getAtom2(i);

      atom.charge = _atoms[i].charge;
      atom.isotope = _atoms[i].isotope;
      atom.implicit_h = _atoms[i].hydrogens; // may receive -1, which means 'unspecified'

      if (query)
         mol.getQueryAtom2(i) = _atoms[i].query;
   }

   for (i = 0; i < _bonds.size(); i++)
   {
      int type = _bonds[i].type;
      int beg = _bonds[i].beg;
      int end = _bonds[i].end;

      if (type == -1)
      {
         if (_atoms[beg].aromatic && _atoms[end].aromatic)
            ;//type = BOND_AROMATIC;
         else
            type = BOND_SINGLE;
      }

      mol.addBond(beg, end, type);

      if (query)
         mol.getQueryBond2(i) = _bonds[i].query;
   }

   mol.findBondsInRings();

   for (i = 0; i < _bonds.size(); i++)
   {
      if (mol.getBond(i).type == -1)
      {
         if (mol.bondInRing(i))
         {
            mol.getBond2(i).type = BOND_AROMATIC;
            mol.registerBond(i, mol.getBond(i));
         }
         else
            mol.getBond2(i).type = BOND_SINGLE;
      }
   }

   _calcStereocenters(mol);
   _calcCisTrans(mol, query);

   while (!_scanner.isEOF())
   {
      char c = _scanner.readChar();

      if (isspace(c))
         continue;
      if (c == '|')
         _readOtherStuff(mol, query);
   }

   if (!inside_rsmiles && !query)
      checkQueryAtoms(mol);

   if (!inside_rsmiles)
   {
     _scanner.skipSpace();
     if (!_scanner.isEOF())
        _scanner.readString(mol.name, true);
   }

   if (reaction_atom_mapping != 0)
   {
      reaction_atom_mapping->clear_resize(mol.vertexCount());
      reaction_atom_mapping->zerofill();
      for (i = 0; i < _atoms.size(); i++)
         reaction_atom_mapping->at(i) = _atoms[i].aam;
   }
}

void SmilesLoader::checkQueryAtoms (const Molecule &mol)
{
   for (int i = 0; i < _atoms.size(); i++)
   {
      if (_atoms[i].query.type != 0 && !mol.isPseudoAtom(i))
         throw Error("'*' atoms accepted only within queries (unless they are pseudo-atoms)");
   }
}


void SmilesLoader::_calcStereocenters (Molecule& mol)
{
   int i, j, tmp;
   
   for (i = 0; i < _atoms.size(); i++)
   {
      if (_atoms[i].chirality > 0)
      {
         QS_DEF(Array<int>, saved);
         MoleculeStereocenters &stereocenters = mol.getStereocenters();

         saved.clear_resize(_atoms.size());
         saved.zerofill();

         int pyramid[4] = {-1, -1, -1, -1};
         int counter = 0;
         int h_index = -1;

         for (j = 0; j < _atoms.size(); j++)
         {
            if (j == i)
            {
               if (mol.getVertex(i).degree() == 3)
               {
                  if (counter >= 4)
                     throw Error("too many bonds for chiral atom %d", i);
                  h_index = counter;
                  pyramid[counter++] = -1;
               }
               for (int k = 0; k < _closures.size(); k++)
               {
                  if (_closures[k].beg == i && !saved[_closures[k].end])
                  {
                     if (counter >= 4)
                        throw Error("too many bonds for chiral atom %d", i);
                     pyramid[counter++] = _closures[k].end;
                     saved[_closures[k].end] = 1;
                  }
               }
            }
            else if (mol.findEdgeIndex(i, j) != -1 && !saved[j])
            {
               if (counter >= 4)
                  throw Error("too many bonds for chiral atom %d", i);
               pyramid[counter++] = j;
               saved[j] = 1;
            }
         }

         if (counter < 3)
            throw Error("only %d bonds for chiral atom %d", counter, i);

         if (counter == 4)
         {
            j = pyramid[0];
            pyramid[0] = pyramid[1];
            pyramid[1] = pyramid[2];
            pyramid[2] = pyramid[3];
            pyramid[3] = j;
            
            if (h_index == 0)
               h_index = 3;
            else if (h_index > 0)
               h_index--;
         }

         if (h_index >= 0)
         {
            if (counter != 4)
               throw Error("implicit hydrogen not allowed with %d neighbor atoms", counter - 1);
            
            bool parity = true;

            for (j = h_index; j < 3; j++)
            {
               __swap(pyramid[j], pyramid[j + 1], tmp);
               parity = !parity;
            }

            if (!parity)
               __swap(pyramid[0], pyramid[1], tmp);
         }

         if (_atoms[i].chirality == 2)
            __swap(pyramid[0], pyramid[1], j);

         stereocenters.add(i, MoleculeStereocenters::ATOM_AND, 1, pyramid);
      }
   }
}

void SmilesLoader::_calcCisTrans (Molecule &mol, bool query)
{
   QS_DEF(Array<int>, dirs);
   int i;

   dirs.clear();

   for (i = 0; i < _bonds.size(); i++)
      dirs.push(_bonds[i].dir);

   // there could be bonds added to stereocenters
   for (; i < mol.edgeEnd(); i++)
      dirs.push(0);

   mol.findBondsInRings();
   mol.cis_trans.buildFromSmiles(mol, dirs.ptr());

   if (query)
      for (i = 0; i < _bonds.size(); i++)
         if (mol.cis_trans.getParity(i) != 0)
            mol.getQueryBond2(i).stereo_care = 1;
}

void SmilesLoader::_readOtherStuff (Molecule &mol, bool query)
{
   MoleculeStereocenters &stereocenters = mol.getStereocenters();
   
   bool highlighting_initialized = false;

   while (1)
   {
      char c = _scanner.readChar();

      if (c == '|')
         break;

      if (c == 'w')
      {
         if (_scanner.readChar() != ':')
            throw Error("colon expected after 'w'");

         while (isdigit(_scanner.lookNext()))
         {
            int idx = _scanner.readUnsigned();

            stereocenters.add(idx, MoleculeStereocenters::ATOM_ANY, 0, false);

            if (_scanner.lookNext() == ',')
               _scanner.skip(1);
         }
      }
      else if (c == 'a')
      {
         if (_scanner.readChar() != ':')
            throw Error("colon expected after 'a'");

         while (isdigit(_scanner.lookNext()))
         {
            int idx = _scanner.readUnsigned();

            stereocenters.setType(idx, MoleculeStereocenters::ATOM_ABS, 0);

            if (_scanner.lookNext() == ',')
               _scanner.skip(1);
         }
      }
      else if (c == 'o')
      {
         int groupno = _scanner.readUnsigned();

         if (_scanner.readChar() != ':')
            throw Error("colon expected after 'o'");
         
         while (isdigit(_scanner.lookNext()))
         {
            int idx = _scanner.readUnsigned();

            stereocenters.setType(idx, MoleculeStereocenters::ATOM_OR, groupno);

            if (_scanner.lookNext() == ',')
               _scanner.skip(1);
         }
      }
      else if (c == '&')
      {
         int groupno = _scanner.readUnsigned();

         if (_scanner.readChar() != ':')
            throw Error("colon expected after '&'");

         while (isdigit(_scanner.lookNext()))
         {
            int idx = _scanner.readUnsigned();

            stereocenters.setType(idx, MoleculeStereocenters::ATOM_AND, groupno);

            if (_scanner.lookNext() == ',')
               _scanner.skip(1);
         }
      }
      else if (c == '^')
      {
         int rad = _scanner.readIntFix(1);
         int radical;

         if (rad == 1)
            radical = RADICAL_DOUPLET;
         else if (rad == 3)
            radical = RADICAL_SINGLET;
         else if (rad == 4)
            radical = RADICAL_TRIPLET;
         else
            throw Error("unsupported radical number: %d", rad);

         if (_scanner.readChar() != ':')
            throw Error("colon expected after radical number");

         while (isdigit(_scanner.lookNext()))
         {
            int idx = _scanner.readUnsigned();

            mol.getAtom2(idx).radical = radical;

            if (_scanner.lookNext() == ',')
               _scanner.skip(1);
         }
      }
      else if (c == '$')
      {
         QS_DEF(Array<char>, label);

         for (int i = 0; i < mol.vertexCount(); i++)
         {
            label.clear();

            while (1)
            {
               if (_scanner.isEOF())
                  throw Error("end of input while reading $...$ block");
               c = _scanner.readChar();
               if (c == ';' || c == '$')
                  break;
               label.push(c);
            }
            if (c == '$' && i != mol.vertexCount() - 1)
               throw Error("only %d atoms found in pseudo-atoms $...$ block", i + 1);
            if (c == ';' && i == mol.vertexCount() - 1)
               throw Error("extra ';' in pseudo-atoms $...$ block");
            if (label.size() > 0)
            {
               label.push(0);
               mol.setPseudoAtom(i, label.ptr());
               if (query)
                  mol.getQueryAtom2(i).type = 0;
            }
         }
      }
      else if (c == 'h')
      {
         c = _scanner.readChar();

         int a = false;

         if (c == 'a')
            a = true;
         else if (c != 'b')
            throw Error("expected 'a' or 'b' after 'h', got '%c'", c);

         if (_scanner.readChar() != ':')
            throw Error("colon expected after 'h%c'", a ? 'a' : 'b');

         if (highlighting != 0 && !highlighting_initialized)
         {
            highlighting->init(mol);
            highlighting_initialized = true;
         }

         while (isdigit(_scanner.lookNext()))
         {
            int idx = _scanner.readUnsigned();

            if (highlighting != 0)
            {
               if (a)
                  highlighting->onVertex(idx);
               else
                  highlighting->onEdge(idx);
            }

            if (_scanner.lookNext() == ',')
               _scanner.skip(1);
         }
      }
   }
}
